import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from keras.models import Sequential
from keras.layers import LSTM, Dense, Input
from keras.callbacks import EarlyStopping


# 1. 生成示例数据（包含多个SPU）
np.random.seed(23)

data_file = "sales30.xlsx"
output_file = "predictions.xlsx"

data = pd.read_excel(data_file)

# 创建DataFrame并使用新字段名
df = pd.DataFrame(data, columns=['spu_code', 'date_id', 'sale_qty'])
spu_codes = df["spu_code"].unique()
df = df.sort_values(['spu_code', 'date_id']).reset_index(drop=True)
# 修复1：确保date_id列是日期时间类型
df['date_id'] = pd.to_datetime(df['date_id'])

# 2. 定义预测函数
def predict_sales(spu_data, days_to_predict=14):
    """预测单个SPU的未来销量"""
    # 数据预处理
    scaler = MinMaxScaler()
    scaled_data = scaler.fit_transform(spu_data[['sale_qty']])

    # 创建时间序列数据集
    def create_dataset(data, time_step=7):
        X, y = [], []
        for i in range(len(data) - time_step):
            X.append(data[i:i + time_step, 0])
            y.append(data[i + time_step, 0])
        return np.array(X), np.array(y)

    time_step = 7
    X, y = create_dataset(scaled_data, time_step)
    X = X.reshape(X.shape[0], time_step, 1)

    # 构建并训练LSTM模型
    model = Sequential([
        Input(shape=(time_step, 1)),
        LSTM(32, return_sequences=True),
        LSTM(16),
        Dense(1)
    ])

    model.compile(optimizer='adam', loss='mse')

    # 使用早停法防止过拟合
    early_stop = EarlyStopping(monitor='loss', patience=10, restore_best_weights=True)
    model.fit(X, y, epochs=100, batch_size=4, callbacks=[early_stop], verbose=0)

    # 预测未来销量
    last_sequence = scaled_data[-time_step:]
    future_predictions = []

    for _ in range(days_to_predict):
        x_input = last_sequence.reshape((1, time_step, 1))
        pred = model.predict(x_input, verbose=0)[0, 0]
        future_predictions.append(pred)
        last_sequence = np.append(last_sequence[1:], [[pred]], axis=0)

    # 反归一化预测结果
    future_predictions = scaler.inverse_transform(np.array(future_predictions).reshape(-1, 1))
    return future_predictions.flatten()


# 3. 为每个SPU预测未来14天销量
results = {}
last_date = df['date_id'].max()
future_dates = [last_date + pd.Timedelta(days=i) for i in range(1, 15)]


for spu in df['spu_code'].unique():
    spu_data = df[df['spu_code'] == spu][['sale_qty']]
    predictions = predict_sales(spu_data)
    results[spu] = predictions

    # # 可视化历史数据和预测
    # plt.figure(figsize=(10, 4))
    # plt.plot(spu_data.index, spu_data['sale_qty'], 'b-o', label='Historical Sale Qty')
    # future_index = range(len(spu_data), len(spu_data) + 14)
    # plt.plot(future_index, predictions, 'r--o', label='Future Prediction')
    # plt.title(f'Sale Qty Forecast for SPU: {spu}')
    # plt.xlabel('Days')
    # plt.ylabel('Sale Quantity')
    # plt.axvline(x=len(spu_data) - 0.5, color='gray', linestyle='--', label='Current Day')
    # plt.legend()
    # plt.grid(True, linestyle='--', alpha=0.7)
    # plt.tight_layout()
    # plt.show()

# 4. 创建预测结果DataFrame
prediction_data = []
for spu in spu_codes:
    for i, date in enumerate(future_dates):
        prediction_data.append({
            'spu_code': spu,
            'date_id': date,
            'predicted_sale_qty': results[spu][i]
        })

result_df = pd.DataFrame(prediction_data)

# 5. 输出预测结果
print("\nFuture Sale Quantity Predictions (Next 14 Days):")
print(result_df)

# 6. 保存结果到CSV
result_df.to_csv('spu_sale_forecast.csv', index=False)
print("\n预测结果已保存到 spu_sale_forecast.csv")

# 7. 可选：按SPU汇总展示预测结果
print("\n预测结果汇总（按SPU分组）:")
for spu in spu_codes:
    spu_df = result_df[result_df['spu_code'] == spu]
    print(f"\nSPU: {spu}")
    print(spu_df[['date_id', 'predicted_sale_qty']].reset_index(drop=True))